import pysam
import argparse
import pod5 as p5
#import numpy as np
from tqdm import tqdm
from pathlib import Path

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--inbam','-i', type=str, default="/public/data/wangjianxin/nipeng/workspace/arab_col0_csu2_nextomics.ont_r10.4.1_simplex.pod5_pass.bam", required=False,
                        help="bam ready to be split")
    parser.add_argument('--out','-o', type=str, default="/public/home/hpc234701005/data/R10.4/arab/train/arab_col0_train.bam", required=False,
                        help="where the readid to store")
    parser.add_argument('--chr', type=str,default="chr1",
                        help="id")
    return parser.parse_args()



def calculate_median(array):
    sorted_array = sorted(array)
    length = len(sorted_array)

    if length % 2 == 0:
        # 如果数组长度为偶数，取中间两个数的平均值
        middle_left = sorted_array[length // 2 - 1]
        middle_right = sorted_array[length // 2]
        median = (middle_left + middle_right) / 2
    else:
        # 如果数组长度为奇数，取中间数
        median = sorted_array[length // 2]

    return median
def calculate_mean(array):
    total = sum(array)
    mean = total / len(array)
    return mean

def deal_bam(inbam,out,process_chr):
    input_bam = pysam.AlignmentFile(inbam, 'rb',check_sq=False)
    wf=open(out, 'w')
    pi_num=0
    query_length=[]
    #ids=set()
    #key_sep = "||"
    for read in tqdm(input_bam,desc="deal bam"):
        reference_name=read.reference_name
        if reference_name is None:
            continue
        if reference_name not in process_chr:
            continue
        if read.has_tag("pi"):
            read_name=read.get_tag("pi")
            pi_num+=1
        else:
            read_name=read.query_name

            
        wf.write(read_name+'\n')
    
    print('have pi read number is {}'.format(pi_num))
    wf.close()


if __name__ == '__main__':
    args=parse_args()
    deal_bam(args.inbam,args.out,args.chr)
